import matplotlib.pyplot as plt

from copy import deepcopy
from random import sample
from random import randint
from random import shuffle

from optimizer.vcmtsp.deprecated.operators import Lifecycle
from optimizer.vcmtsp.deprecated.operators import minimal_insertion
from optimizer.vcmtsp.deprecated.operators import two_opt, jaccard


# plt.figure(figsize=(8, 4))

class GA:
    """
    Genetic algorithm (GA) on VCMTSP
    """
    def __init__(self, data, NP: int, Pc: float, Pm: float, method):
        # -- parameters -- #
        self.life = Lifecycle()
        self.__data = data
        self.__n, self.__m = self.__data.n, self.__data.m
        self.__distances = self.__data.distances
        self.__A, self.__S, self.__C = self.__data.A, self.__data.S, self.__data.C
        self.__NP, self.__Pc, self.__Pm = NP, Pc, Pm

        self.method = method

        self.__population = []
        for i in range(self.__NP):
            routes = self.__data.init_routes()
            length = self.__data.cost(routes=routes)
            self.__population.append([routes, length])
        self.gBest = deepcopy(min(self.__population, key=lambda _: _[1]))
        self.life.update(gBest=self.gBest[1])

    def optimize(self, iterations=None, evaluations=None, deadline=None):
        self.life.set_stopping_condition(iterations, evaluations, deadline)
        while self.life.check_stopping_condition(mute=False) is False:
            self.__crossover()

            if self.method == '1':
                self.__outer_mutation_each_city()
            elif self.method == '2':
                self.__outer_mutation()
            else:
                print(f'Error. method={self.method}')
                quit(-1)
            self.__inner_mutation()
            self.__top_NP_selection()

            if self.gBest[1] > min(self.__population, key=lambda _: _[1])[1]:
                self.gBest = deepcopy(min(self.__population, key=lambda _: _[1]))
            self.life.update(it=1, gBest=self.gBest[1])

        self.life.update(data=self.gBest[0])
        return self.life

    def __crossover(self):
        offspring = []
        shuffle(self.__population)
        for i in range(self.__NP):
            if i == 0:
                a, b = self.__population[-1], self.__population[0]
            else:
                a, b = self.__population[i - 1], self.__population[i]
            c = self.__inner_cross(a, b)  # produce a child
            print(a[0])
            print(b[0])
            plt.subplot(1, 2, 1)
            self.__data.draw(routes=a[0], emphasize=1, title='pA')
            plt.subplot(1, 2, 2)
            self.__data.draw(routes=b[0], emphasize=1, title='pB')
            plt.show()
            offspring.append(c)
        self.__population += offspring  # merge offspring with the former population

    def __crossover_pairwise(self):
        offspring = []
        shuffle(self.__population)  # shuffle population
        for i in range(self.__NP):
            if i % 2 == 1:  # cross each pair of individuals
                a, b = self.__population[i - 1], self.__population[i]
                c = self.__inner_cross(a, b)  # produce a child
                offspring.append(c)
                c = self.__inner_cross(b, a)  # produce a child
                offspring.append(c)
        self.__population += offspring  # merge offspring with the former population

    def __inner_cross(self, a, b):
        child = []  # child routes are produced based on a while referring to b
        removed = []  # removed cities
        length = a[1]  # total length of routes in child
        for a_route in a[0]:
            # select the most similar route to cross
            sim = []
            for b_route in b[0]:
                sim.append(jaccard(a_route, b_route))
            k = sim.index(max(sim))
            b_route = b[0][k]
            route = a_route[:]
            for i in range(len(route) - 1, -1, -1):  # traverse in the inverted order
                if len(route) > 3:
                    if route[i] not in b_route:  # delete the cities that not in b
                        length += self.__distances[route[i - 1]][route[i + 1]] \
                                  - self.__distances[route[i - 1]][route[i]] \
                                  - self.__distances[route[i]][route[i + 1]]
                        removed.append(route[i])  # store the deleted cities
                        route.remove(route[i])

            child.append(route)

        # re-insert the removed cities by minimal cost
        shuffle(removed)
        child, cost = minimal_insertion(distance_func=self.__distances, data_S=self.__S,
                                        cities=removed, routes=child)
        length += cost
        return [child, length]

    def __outer_mutation(self):
        for _, solution in enumerate(self.__population):  # for each individual
            k1, k2 = sample(range(0, self.__m), 2)  # random select 2 routes
            cities = set(solution[0][k1]) & set(self.__C[k2 + 1])  # cities which are in k1, and also accessible to k2
            if len(cities) > 1 and len(solution[0][k1]) > 3:
                city = sample(cities, 1)[0]  # random select one shared city
                p, q = solution[0][k1].index(city) - 1, solution[0][k1].index(city) + 1
                p, q = solution[0][k1][p], solution[0][k1][q]  # sequence p-city-q
                solution[0][k1].remove(city)  # remove city
                solution[1] += self.__distances[p][q] - self.__distances[p][city] - self.__distances[city][q]
                # find the best position in route k2 to insert city
                _i, min_diff = -1, 1e99
                for i in range(1, len(solution[0][k2])):
                    p, q = solution[0][k2][i - 1], solution[0][k2][i]
                    diff = self.__distances[p][city] + self.__distances[city][q] - self.__distances[p][q]
                    if min_diff > diff:
                        min_diff = diff
                        _i = i
                solution[0][k2].insert(_i, city)
                solution[1] += min_diff

    def __inner_mutation(self):
        for p in range(self.__NP):  # for each individual
            for k in range(self.__m):  # do 2-opt on each route
                self.__population[p][0][k], improve = two_opt(self.__population[p][0][k], self.__distances)
                self.__population[p][1] += improve

    def __top_NP_selection(self):
        self.__population = sorted(self.__population, key=lambda _: _[1])
        self.__population = self.__population[:self.__NP]

    def __tournament_selection(self):
        winners = []
        for i in range(self.__NP):
            a, b = sample(self.__population, k=2)
            if a[1] < b[1]:
                winners.append(a)
                self.__population.remove(a)
            else:
                winners.append(b)
                self.__population.remove(b)
        self.__population = winners

    # =============== deprecated *******
    def __outer_mutation_dep(self, re_loc=False, swap=False):
        for _, solution in enumerate(self.__population):  # for each individual
            if re_loc:
                k1, k2 = sample(range(0, self.__m), 2)  # random select 2 routes
                assert k1 != k2
                for i, a in enumerate(solution[0][k1]):
                    if self.__A[a][k2 + 1]:  # if city a is accessible to k2
                        p, q = solution[0][k1][i - 1], solution[0][k1][i + 1]
                        remove_a = self.__distances[p][q] - self.__distances[p][a] - self.__distances[a][q]
                        _j, min_diff = -1, 1e99
                        for j in range(1, len(solution[0][k2]) - 1):  # find minimal insertion
                            m, b = solution[0][k2][j - 1], solution[0][k2][j]
                            diff = self.__distances[m][a] + self.__distances[a][b] - self.__distances[m][b]
                            diff += remove_a
                            if min_diff > diff:
                                min_diff = diff
                                _j = j
                        if min_diff < -1e-10 and len(solution[0][k1]) > 3:
                            self.__population[_][0][k1].remove(a)
                            self.__population[_][0][k2].insert(_j, a)
                            self.__population[_][1] += min_diff

                for i, a in enumerate(solution[0][k2]):
                    if self.__A[a][k1 + 1]:  # if city a is accessible to k1
                        p, q = solution[0][k2][i - 1], solution[0][k2][i + 1]
                        remove_a = self.__distances[p][q] - self.__distances[p][a] - self.__distances[a][q]
                        _j, min_diff = -1, 1e99
                        for j in range(1, len(solution[0][k1]) - 1):  # find minimal insertion
                            m, b = solution[0][k1][j - 1], solution[0][k1][j]
                            diff = self.__distances[m][a] + self.__distances[a][b] - self.__distances[m][b]
                            diff += remove_a
                            if min_diff > diff:
                                min_diff = diff
                                _j = j
                        if min_diff < -1e-10 and len(solution[0][k2]) > 3:
                            self.__population[_][0][k2].remove(a)
                            self.__population[_][0][k1].insert(_j, a)
                            # self.__population[_][1] += min_diff
                            new_cost = self.__data.cost(routes=self.__population[_][0])
                            if abs(min_diff - (new_cost - self.__population[_][1])) > 1e-6:
                                print(f'error: min_diff={min_diff}, cost_diff={(new_cost - self.__population[_][1])}')
                            self.__population[_][1] = new_cost

            if swap:
                k1, k2 = sample(range(0, self.__m), 2)  # random select 2 routes
                assert k1 != k2
                for i, a in enumerate(solution[0][k1]):
                    if a != self.__data.depot and self.__data.A[a][k2 + 1] == 1:  # if city a is accessible to k2
                        p, q = solution[0][k1][i - 1], solution[0][k1][i + 1]
                        _j, min_diff = -1, 1e99
                        for j, b in enumerate(solution[0][k2]):  # find minimal swap
                            if b != self.__data.depot and self.__data.A[b][k1 + 1] == 1:  # if b is accessible to k1
                                m, n = solution[0][k2][j - 1], solution[0][k2][j + 1]
                                diff = self.__distances[m][a] + self.__distances[a][n] \
                                       + self.__distances[p][b] + self.__distances[p][q] \
                                       - self.__distances[p][a] - self.__distances[a][q] \
                                       - self.__distances[m][b] - self.__distances[b][n]
                                if min_diff > diff:
                                    min_diff = diff
                                    _j = j
                        if min_diff < -1e-10:
                            self.__population[_][0][k1][i], self.__population[_][0][k2][_j] = \
                                self.__population[_][0][k2][_j], self.__population[_][0][k1][i]
                            # self.__population[_][1] += min_diff
                            new_cost = self.__data.cost(routes=self.__population[_][0])
                            if abs(min_diff - (new_cost - self.__population[_][1])) > 1e-6:
                                print(f'error: min_diff={min_diff}, cost_diff={(new_cost - self.__population[_][1])}')
                            self.__population[_][1] = new_cost

    def __mutation(self, swap=False, re_loc=False):
        for _, solution in enumerate(self.__population):  # each individual
            k = randint(0, self.__m - 1)  # random select a route
            i = randint(1, len(solution[0][k]) - 2)  # random select a city except the depot
            p, a, q = solution[0][k][i - 1], solution[0][k][i], solution[0][k][i + 1]  # prev, curr, next
            S = self.__data.S[a]  # accessible routes (salesmen) of the city 'a'
            _s, _j = -1, -1  # swap/insert position
            # ======== remove & re-insert =========
            if re_loc:
                min_diff = 1e99
                for s in S:  # search in each accessible route
                    for j in range(1, len(solution[0][s - 1]) - 1):
                        m, n = solution[0][s - 1][j - 1], solution[0][s - 1][j]
                        if a == m or a == n:
                            continue
                        diff = self.__distances[m][a] + self.__distances[a][n] - self.__distances[m][n] \
                               + self.__distances[p][q] - self.__distances[p][a] - self.__distances[a][q]
                        if min_diff > diff:
                            min_diff = diff
                            _s, _j = s, j  # find the minimal swap position
                if min_diff > 0:
                    continue
                self.__population[_][0][k].remove(a)  # remove a from route k
                self.__population[_][0][_s - 1].insert(_j, a)  # insert a to route s-1 at position _j
                self.__population[_][1] += min_diff  # adjust traveling cost
            # =============== swap ===============
            if swap:
                min_diff = 1e99
                for s in S:  # search in each accessible route
                    for j in range(1, len(solution[0][s - 1]) - 1):
                        b = solution[0][s - 1][j]  # city 'b'
                        if (k + 1) in self.__data.S[b]:  # check if b can be swapped
                            m, n = solution[0][s - 1][j - 1], solution[0][s - 1][j + 1]  # prev, curr, next
                            diff = self.__distances[m][a] + self.__distances[a][n] \
                                   - self.__distances[m][b] - self.__distances[b][n] \
                                   + self.__distances[p][b] + self.__distances[b][q] \
                                   - self.__distances[p][a] - self.__distances[a][q]
                            if min_diff > diff:
                                min_diff = diff
                                _s, _j = s, j  # find the minimal swap position
                    # print(k, i)
                    print(_s, _j, min_diff)
                self.__population[p][0][k][i], self.__population[p][0][_s - 1][_j] = \
                    self.__population[p][0][_s - 1][_j], self.__population[p][0][k][i]  # swap a and b

                self.__population[p][1] += min_diff  # adjust traveling cost

    def __outer_mutation_each_city(self):
        """
        perform on each city
        @return:
        """
        for _, solution in enumerate(self.__population):  # for each individual
            k1, k2 = sample(range(0, self.__m), 2)  # random select 2 routes
            for i in range(len(solution[0][k1]) - 2, 0, -1):
                a = solution[0][k1][i]
                if self.__A[a][k2 + 1]:  # if city a is accessible to k2
                    p, q = solution[0][k1][i - 1], solution[0][k1][i + 1]
                    remove_a = self.__distances[p][q] - self.__distances[p][a] - self.__distances[a][q]
                    _j, min_diff = -1, 1e99
                    for j in range(1, len(solution[0][k2]) - 1):  # find minimal insertion
                        m, b = solution[0][k2][j - 1], solution[0][k2][j]
                        diff = self.__distances[m][a] + self.__distances[a][b] - self.__distances[m][b]
                        diff += remove_a
                        if min_diff > diff:
                            min_diff = diff
                            _j = j
                    if min_diff < -1e-10 and len(solution[0][k1]) > 3:
                        self.__population[_][0][k1].remove(a)
                        self.__population[_][0][k2].insert(_j, a)
                        new_cost = self.__data.cost(routes=self.__population[_][0])
                        diff = new_cost - self.__population[_][1]
                        if abs(min_diff - diff) > 1e-6:
                            print(f'error: min_diff={min_diff}, cost_diff={diff}')
                        self.__population[_][1] += min_diff

            for i in range(len(solution[0][k2]) - 2, 0, -1):
                a = solution[0][k2][i]
                if self.__A[a][k1 + 1]:  # if city a is accessible to k1
                    p, q = solution[0][k2][i - 1], solution[0][k2][i + 1]
                    remove_a = self.__distances[p][q] - self.__distances[p][a] - self.__distances[a][q]
                    _j, min_diff = -1, 1e99
                    for j in range(1, len(solution[0][k1]) - 1):  # find minimal insertion
                        m, b = solution[0][k1][j - 1], solution[0][k1][j]
                        diff = self.__distances[m][a] + self.__distances[a][b] - self.__distances[m][b]
                        diff += remove_a
                        if min_diff > diff:
                            min_diff = diff
                            _j = j
                    if min_diff < -1e-10 and len(solution[0][k2]) > 3:
                        self.__population[_][0][k2].remove(a)
                        self.__population[_][0][k1].insert(_j, a)
                        new_cost = self.__data.cost(routes=self.__population[_][0])
                        diff = new_cost - self.__population[_][1]
                        if abs(min_diff - diff) > 1e-6:
                            print(f'error: min_diff={min_diff}, cost_diff={diff}')
                        self.__population[_][1] += min_diff
